#include <iostream>
#include <vector>
#include <random>
#include <stdexcept>
#include <algorithm>
#include <openssl/sha.h>
#include <gmp.h>
#include <gmpxx.h>

using namespace std;

// 全局参数: 素数 p = 2^128 + 51
const mpz_class P = (mpz_class(1) << 128) + 51;

// 辅助函数：将mpz_class转换为字节向量
vector<uint8_t> mpz_to_bytes(const mpz_class& num, size_t len = 16) {
    vector<uint8_t> bytes(len, 0);
    mpz_class n = num;
    
    // 导出大整数到字节数组
    size_t count;
    mpz_export(bytes.data(), &count, 1, 1, 0, 0, n.get_mpz_t());
    
    // 如果需要，反转字节顺序为大端序
    // if (count < len) {
    //     reverse(bytes.begin(), bytes.begin() + count);
    // }
    
    return bytes;
}

// 辅助函数：将字节向量转换为mpz_class
mpz_class bytes_to_mpz(const vector<uint8_t>& bytes) {
    mpz_class result;
    // 导入字节数组到大整数
    mpz_import(result.get_mpz_t(), bytes.size(), 1, 1, 0, 0, bytes.data());
    return result;
}

// 模逆元计算
mpz_class mod_inv(const mpz_class& a, const mpz_class& m) {
    mpz_class inv;
    if (mpz_invert(inv.get_mpz_t(), a.get_mpz_t(), m.get_mpz_t()) == 0) {
        throw runtime_error("mod_inv does not exist!");
    }
    return inv;
}

// 密钥派生函数 (基于SHA-256)
vector<uint8_t> kdf(const mpz_class& key, size_t len) {
    vector<uint8_t> stream;
    // 将主密钥转换为固定长度字节（16字节）
    vector<uint8_t> key_bytes = mpz_to_bytes(key, 16);
    uint32_t ctr = 0; // 计数器

    while (stream.size() < len) {
        // 构造输入：密钥字节 + 计数器（4字节，大端序）
        vector<uint8_t> input = key_bytes;
        for (int i = 3; i >= 0; --i) {
            input.push_back(static_cast<uint8_t>((ctr >> (i*8)) & 0xFF));
        }

        // 计算SHA-256
        uint8_t hash[SHA256_DIGEST_LENGTH];
        SHA256(input.data(), input.size(), hash);
        
        // 将哈希值添加到流
        size_t needed = len - stream.size();
        size_t n = min(static_cast<size_t>(SHA256_DIGEST_LENGTH), needed);
        stream.insert(stream.end(), hash, hash + n);

        // 更新计数器
        ctr++;
    }
    return stream;
}

// 加密函数
pair<vector<uint8_t>, vector<pair<int, mpz_class>>> 
encrypt(const string& plaintext) {
    // 初始化GMP随机状态
    gmp_randstate_t state;
    gmp_randinit_default(state);
    
    // 使用系统随机设备作为种子
    random_device rd;
    gmp_randseed_ui(state, rd());

    // 生成128位主密钥 (0 到 2^128-1)
    mpz_class K;
    mpz_urandomb(K.get_mpz_t(), state, 128);
    
    // 随机多项式系数
    mpz_class a1, a2;
    mpz_urandomm(a1.get_mpz_t(), state, P.get_mpz_t());
    mpz_urandomm(a2.get_mpz_t(), state, P.get_mpz_t());
    
    // 清理随机状态
    gmp_randclear(state);

    // 计算密钥份额 (x=1,2,3,4)
    vector<pair<int, mpz_class>> shares;
    for (int x = 1; x <= 4; ++x) {
        mpz_class fx = (K + a1*x + a2*x*x) % P;
        if (fx < 0) fx += P;
        shares.push_back({x, fx});
    }
    
    // 加密明文
    vector<uint8_t> key_stream = kdf(K, plaintext.size());
    vector<uint8_t> ciphertext;
    for (size_t i = 0; i < plaintext.size(); ++i) {
        ciphertext.push_back(plaintext[i] ^ key_stream[i]);
    }
    
    return {ciphertext, shares};
}

// 解密函数
string decrypt(const vector<uint8_t>& ciphertext, 
               const vector<pair<int, mpz_class>>& shares, const int limit_num) {
    // if (shares.size() < 3) {
    //     throw runtime_error("At least three secret shares are needed!");
    // }
    
    int x[limit_num];
    mpz_class y[limit_num];
    mpz_class denom[limit_num];
    mpz_class L[limit_num];
    // 拉格朗日插值恢复K
    // const auto& [x1, y1] = shares[0];
    // const auto& [x2, y2] = shares[1];
    // const auto& [x3, y3] = shares[2];
    for(int i = 0; i < limit_num; i++){
        x[i] = shares[i].first;
        y[i] = shares[i].second;
    }
    
    // 计算分母
    // mpz_class denom1 = (x1 - x2) * (x1 - x3);
    // mpz_class denom2 = (x2 - x1) * (x2 - x3);
    // mpz_class denom3 = (x3 - x1) * (x3 - x2);
    for(int i = 0; i < limit_num; i++){
        denom[i] = 1;
        for(int j = 0; j < limit_num; j++){
            if(j != i)
                denom[i] = denom[i]*(x[i] - x[j]);
        }
    }
    
    // // 计算基函数在0处的值
    // mpz_class L1 = (x2 * x3) * mod_inv(denom1, P) % P;
    // mpz_class L2 = (x1 * x3) * mod_inv(denom2, P) % P;
    // mpz_class L3 = (x1 * x2) * mod_inv(denom3, P) % P;
    for(int i = 0; i < limit_num; i++){
        L[i] = 1;
        for(int j = 0; j < limit_num; j++){
            if(j != i)
                L[i] = L[i]*(-x[j]);
        }
        L[i] = L[i] * mod_inv(denom[i], P) % P;
    }
    
    // mpz_class K = (y1*L1 + y2*L2 + y3*L3) % P;
    mpz_class K = 0;
    for(int i = 0; i < limit_num; i++){
            K += y[i] * L[i];
    }
    K = K % P;
    if (K < 0) K += P;
    
    // 解密密文
    vector<uint8_t> key_stream = kdf(K, ciphertext.size());
    string plaintext;
    for (size_t i = 0; i < ciphertext.size(); ++i) {
        plaintext.push_back(ciphertext[i] ^ key_stream[i]);
    }
    return plaintext;
}

// 打印十六进制数据
void print_hex(const vector<uint8_t>& data) {
    for (uint8_t c : data) {
        printf("%02X", c);
    }
    cout << endl;
}

// 打印密钥份额
void print_shares(const vector<pair<int, mpz_class>>& shares) {
    for (const auto& [x, y] : shares) {
        cout << "Share " << x << ": ";
        vector<uint8_t> bytes = mpz_to_bytes(y);
        print_hex(bytes);
    }
}

int main() {
    int limit_num;
    cout << "Limit: ";
    cin >> limit_num;
    getchar();
    int x[limit_num], ciphertext_len;

    // vector<uint8_t> ciphertext={0x22,0x92,0x68,0x01,0x72,0xC6,0x54,0x2B,0x9A,0x95,0x7A,0xED,0xED,0x7B,0xF9,0x1F,0xB5,0x7C,0xEA,0x11,0xE2,0x4A,0x3E,0xCF,0xEA,0xB0,0x0D,0xAB,0xF7,0x91,0xB2,0x53,0x59,0x6C,0xE7,0xFD,0x45,0x48,0x14,0xCE,0x0D,0x46,0x58,0x36,0x6A,0xAE,0x8D,0xF0,0x07,0x15,0x1A,0x81,0x4C,0x28,0x18,0x60,0x31,0x6A,0xB3,0xB9};
    // vector<uint8_t> share1={0x45,0x10,0x23,0xEC,0x8E,0xD6,0x97,0xA2,0x44,0x2C,0x98,0x6E,0xDE,0x43,0x1F,0x6D};
    // vector<uint8_t> share2={0x6F,0xC4,0x3D,0xC7,0x7B,0xCD,0x2E,0xD4,0x2B,0xCA,0xBE,0x46,0x13,0xCA,0x91,0x74};
    // vector<uint8_t> share3={0x10,0xC8,0x6E,0x58,0x27,0xA3,0xB3,0x52,0x56,0x32,0x4A,0x12,0x59,0x25,0xAC,0xD7};

    mpz_t ciphertext_t, shares_t;
    mpz_class ciphertext, shares[limit_num];
    vector<pair<int, mpz_class>> sub_shares;
    string input_cache;

    // mpz_init(ciphertext_t);
    // mpz_init(shares_t);

    cout << "Ciphertext: ";
    getline(cin, input_cache);
    ciphertext.set_str(input_cache, 16);
    ciphertext_len = input_cache.length() / 2;
    for(int i = 0 ; i < limit_num; i++){
        cout << "Share " << i + 1 << " info" << endl << "x" << i + 1 << ": ";
        cin >> x[i];
        cout << "Fx" << i + 1 << ": ";
        getchar();
        getline(cin, input_cache);
        shares[i].set_str(input_cache, 16);
        sub_shares.push_back({x[i], shares[i]});
    }

    vector<uint8_t> ciphertext_bytes = mpz_to_bytes(ciphertext, ciphertext_len);
    print_hex(ciphertext_bytes);
    print_shares(sub_shares);

    string decrypted = decrypt(mpz_to_bytes(ciphertext, ciphertext_len), sub_shares, limit_num);
    cout << "\nDecryption Result: " << decrypted << endl;
    
    return 0;
}
